from anythingllm import AnythingLLM
from deepseek import *
import knowledge
import requests
import json
import os
import time

ANYTHINGLLM_API_KEY = "7QN2S8T-DF1M5VJ-MRFC3VJ-E4HVFG7"
DEEPSEEK_API_KEY = "sk-caf3a3eb52a34c62a38c70b61f5911b4"


chat = DeepSeek.get_cloud_deepseek(DEEPSEEK_API_KEY)
# chat = DeepSeek.get_cloud_deepseek(url='第三方DeepSeek API url', Module='第三方DeepSeek API Module', APIkey='第三方DeepSeek API Key')

llm = AnythingLLM(ANYTHINGLLM_API_KEY)

# 展示所有的工作区
def show_workspaces():
    print("正在获取工作区列表...")
    workspaces = llm.get_spaces()

    os.system('cls')
    if not workspaces:
        print("没有可用的工作区")
        return

    for i, workspace in enumerate(workspaces['workspaces']):
        print(f"{i+1}. {workspace['slug']}")

    workspace_id = input("请输入工作区ID（输入其他内容创建新工作区）：")

    if workspace_id.isdigit():
        workspace_id = int(workspace_id) - 1
        if workspace_id >= 0 or workspace_id < len(workspaces['workspaces']):
            slug = workspaces['workspaces'][workspace_id]['slug']
            return slug

    return llm.creat_space(workspace_id)['workspace']['slug']

def test_deepseek():
    url = "https://api.deepseek.com/user/balance"

    payload={}
    headers = {
    'Accept': 'application/json',
    'Authorization': f'Bearer {DEEPSEEK_API_KEY}'
    }

    response = requests.request("GET", url, headers=headers, data=payload)
    if response.status_code == 200:
        return True
    else:
        return False

def refresh(workspace_slug):
    os.system('cls')
    print('正在处于工作区：', workspace_slug)

def work(workspace_slug):
    refresh(workspace_slug)
    print('正在连接到DeepSeek服务器...')

    if(not test_deepseek()):
        refresh(workspace_slug)
        print('连接失败！')
        return
    
    refresh(workspace_slug)
    print('连接成功！')

    time.sleep(0.5)
    refresh(workspace_slug)

    topic = input('请输入需要生成的知识库关键词：')

    refresh(workspace_slug)

    print('正在生成知识库目录...')

    index = knowledge.get_index(chat, topic)
    knowledge.write_index(index)
    knowledge.write_index(index)
    refresh(workspace_slug)

    print('正在构建知识库...')
    for i in index['list']:
        print(i['name'])
        knowledge.add_first_level_item(i['name'])
        for j in i['list']:
            knowledge.add_second_level_item(j['name'])
            print(f"  - {j['name']}")
            for k in j['list']:
                print(f"    - {k['name']}", end='\t')
                knowledge.add_knowledge(chat, index['name'], i['name'], j['name'], k['name'])
                print('完成')

    refresh(workspace_slug)
    print('构建知识库完成！')

    time.sleep(0.5)
    refresh(workspace_slug)
    print('正在上传知识库到工作区...')

    llm.automate_knowledge_upload(workspace_slug, 'knowledge.md')

    refresh(workspace_slug)
    print('添加知识库成功，可以打开AnythinLLM查看！')

    os.system('pause')


if __name__ == '__main__':
    workspace_slug = show_workspaces()
    work(workspace_slug)
    
    
